import hashlib
import os
from typing import Union # Added Union import
from Crypto.Cipher import AES
from Crypto.PublicKey import RSA
from Crypto.Random import get_random_bytes

# Removed Dify specific imports for redis and storage
# from extensions.ext_redis import redis_client
# from extensions.ext_storage import storage

# Assuming gmpy2_pkcs10aep_cipher is a local utility or a standard library part.
# If it's from Dify's libs, ensure it's copied and doesn't have further Dify dependencies.
from libs import gmpy2_pkcs10aep_cipher

# Define a base path for storing private keys within the application structure
PRIVKEYS_BASE_DIR = os.path.join(os.path.dirname(os.path.dirname(__file__)), "privkeys") # Assuming rsa.py is in backend/libs/

if not os.path.exists(PRIVKEYS_BASE_DIR):
    os.makedirs(PRIVKEYS_BASE_DIR, exist_ok=True)

def generate_key_pair(tenant_id: str):
    private_key = RSA.generate(2048)
    public_key = private_key.publickey()

    pem_private = private_key.export_key()
    pem_public = public_key.export_key()

    # Store private key locally
    tenant_key_dir = os.path.join(PRIVKEYS_BASE_DIR, str(tenant_id))
    if not os.path.exists(tenant_key_dir):
        os.makedirs(tenant_key_dir, exist_ok=True)
    
    filepath = os.path.join(tenant_key_dir, "private.pem")

    with open(filepath, "wb") as f:
        f.write(pem_private)

    return pem_public.decode()


prefix_hybrid = b"HYBRID:"


def encrypt(text: str, public_key: Union[str, bytes]):
    if isinstance(public_key, str):
        public_key = public_key.encode()

    aes_key = get_random_bytes(16)
    cipher_aes = AES.new(aes_key, AES.MODE_EAX)

    ciphertext, tag = cipher_aes.encrypt_and_digest(text.encode())

    rsa_key_obj = RSA.import_key(public_key)
    cipher_rsa = gmpy2_pkcs10aep_cipher.new(rsa_key_obj)

    enc_aes_key = cipher_rsa.encrypt(aes_key)

    encrypted_data = enc_aes_key + cipher_aes.nonce + tag + ciphertext

    return prefix_hybrid + encrypted_data


# Simple in-memory cache for private keys to reduce disk I/O for frequent decryptions
# For a production system, a more robust caching mechanism might be needed if performance is critical.
_private_key_cache = {}

def get_decrypt_decoding(tenant_id: str):
    filepath = os.path.join(PRIVKEYS_BASE_DIR, str(tenant_id), "private.pem")

    # Check cache first
    if filepath in _private_key_cache:
        private_key_pem = _private_key_cache[filepath]
    else:
        try:
            with open(filepath, "rb") as f:
                private_key_pem = f.read()
            _private_key_cache[filepath] = private_key_pem # Cache it
        except FileNotFoundError:
            raise PrivkeyNotFoundError(f"Private key not found for tenant_id: {tenant_id} at {filepath}")

    rsa_key_obj = RSA.import_key(private_key_pem)
    cipher_rsa = gmpy2_pkcs10aep_cipher.new(rsa_key_obj)

    return rsa_key_obj, cipher_rsa


def decrypt_token_with_decoding(encrypted_text: bytes, rsa_key_obj: RSA.RsaKey, cipher_rsa):
    if encrypted_text.startswith(prefix_hybrid):
        encrypted_text = encrypted_text[len(prefix_hybrid) :]

        key_size_bytes = rsa_key_obj.size_in_bytes()
        enc_aes_key = encrypted_text[:key_size_bytes]
        nonce = encrypted_text[key_size_bytes : key_size_bytes + 16]
        tag = encrypted_text[key_size_bytes + 16 : key_size_bytes + 32]
        ciphertext = encrypted_text[key_size_bytes + 32 :]

        aes_key = cipher_rsa.decrypt(enc_aes_key)

        cipher_aes = AES.new(aes_key, AES.MODE_EAX, nonce=nonce)
        decrypted_text = cipher_aes.decrypt_and_verify(ciphertext, tag)
    else:
        decrypted_text = cipher_rsa.decrypt(encrypted_text)

    return decrypted_text.decode()


def decrypt(encrypted_text: bytes, tenant_id: str):
    rsa_key_obj, cipher_rsa = get_decrypt_decoding(tenant_id)
    return decrypt_token_with_decoding(encrypted_text, rsa_key_obj, cipher_rsa)


class PrivkeyNotFoundError(Exception):
    pass

